﻿using Ace;
using Ace.Files;
using Ace.Files.Json;
using Ace.Wpf.Mvvm;
using MalwarePatch.Windows.Check;
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Net;
using System.Net.Mime;
using System.Security.Cryptography;
using System.Security.Cryptography.X509Certificates;
using System.Text;
using System.Threading;
using System.Threading.Tasks;
using System.Windows.Shell;

namespace MalwarePatch.CliTools
{
  class MalwareCheckException : Exception
  {
    public MalwareCheckException(string message) : base(message) { }
    public MalwareCheckException() : base() { }
  }
  class MalwareChecker
  {
    private const string InfoPath = @"assets/malware-remote-info.json";
    private const string UserAgent = @"Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/84.0.4147.105 Safari/537.36";
    private readonly JsonObject infoObject;
    private readonly StringBuilder stringBuilder = new StringBuilder();
    private readonly bool clean = false;
    private static readonly string CheckerOutput = "samples";
    private static readonly string LogFileName = Path.Combine(CheckerOutput, "mwp-check.log");
    private double downloadedCount = 0;

    public MalwareChecker(bool clean = false)
    {
      this.clean = clean;
      infoObject = JsonObject.Parse(AppResource.GetResourceData(InfoPath).ToString(true));
    }
    public static bool TryParse(string[] args)
    {
      if (args.Length == 0)
      {
        return false;
      }
      if (args[0].EqualsIgnoreCase("--malware-check"))
      {
        new CheckWindow(args.Length > 1 && args[1] == "--clean").ShowDialog();
        return true;
      }
      else
      {
        return false;
      }
    }
    public MalwareCheckerModel Model { get; } = new MalwareCheckerModel();
    public bool Debug { get; set; } = false;
    public Task Start()
    {
      stringBuilder.Clear();
      downloadedCount = 0;
      Model.ProgressState = TaskbarItemProgressState.Normal;
      log("Malware checker started");
      return Task.Run(() =>
      {
        if (!Directory.Exists(CheckerOutput))
        {
          Directory.CreateDirectory(CheckerOutput);
        }
        Parallel.ForEach(infoObject, property => downloadCategory(property));
        Stop();
      });
    }
    public void Stop()
    {
      log("Malware checker completed");
      Model.ProgressState = TaskbarItemProgressState.None;
      lock (stringBuilder)
      {
        File.WriteAllText(LogFileName, stringBuilder.ToString());
      }
    }
    private void log(string message)
    {
      lock (stringBuilder)
      {
        stringBuilder
          .Append($"[{DateTime.Now}] ")
          .Append(message)
          .Append(Environment.NewLine);
        System.Diagnostics.Debug.WriteLine(message);
      }
    }
    private void error(string message)
    {
      lock (stringBuilder)
      {
        stringBuilder
          .Append(Environment.NewLine)
          .Append($"[{DateTime.Now}] ")
          .Append(message)
          .Append(Environment.NewLine)
          .Append(Environment.NewLine);
        System.Diagnostics.Debug.WriteLine(message);
      }
    }
    private void downloadCategory(JsonProperty category)
    {
      try
      {
        Model.Message = $"Downloading: {category.Name}";
        downloadFixedUrl(category.ObjectValue["fixed"].ArrayValue);
        // log($"Downloaded fixed urls of {category.Name}");
        downloadRegexUrl(category.ObjectValue["match"].ArrayValue);
        // log($"Downloaded regex urls of {category.Name}");
        downloadedCount++;
        Model.ProgressValue = downloadedCount / infoObject.Count;
      }
      catch (Exception ex)
      when (ex is MalwareCheckException || ex is WebException /*|| ex is CryptographicException*/)
      {
        error($"Error downloading {category.Name}: {ex.Message}");
      }
    }
    private void downloadUrl(string url)
    {
      if (!Debug)
      {
        using (var webClient = new WebClient())
        {
          if (url.StartsWith("//"))
          {
            url = "http:" + url;
          }
          webClient.Headers["user-agent"] = UserAgent;
          var filename = Path.Combine(CheckerOutput, md5(url) + ".exe");
          if (!File.Exists(filename))
          {
            var data = webClient.DownloadData(url);
            File.WriteAllBytes(filename, data);
            log($"Downloaded: {url}");
          }
          else
          {
            log($"Already downloaded: {url}");
          }
          verifyCertificate(filename);
        }
      }
      else
      {
        // Thread.Sleep(3000);
        log($"Debug download: {url}");
      }
    }
    private void verifyCertificate(string path)
    {
      var disallowed = false;
      try
      {
        var certificate = new CertificateFile(new X509Certificate2(path));
        disallowed = certificate.IsDisallowed;
        certificate.Dispose();
      }
      catch (CryptographicException ex)
      {
        error($"Certificate error: {ex.Message}, path={path}");
      }

      if (disallowed)
      {
        if (clean)
        {
          try
          {
            File.Delete(path);
            log($"Verified {path}: Disallowed and deleted.");
          }
          catch (Exception ex)
          {
            error($"Error deleting file {path}: {ex.Message}");
          }
        }
        else
        {
          log($"Verified {path}: Disallowed.");
        }
      }
      else
      {
        log($"Verified {path}: NOT disallowed.");
      }
    }
    private void downloadFixedUrl(JsonArray array)
    {
      array.ToStringList().ForEach(url => downloadUrl(url));
    }
    private void downloadRegexUrl(JsonArray array)
    {
      foreach (var matchInfo in array.ToObjectList())
      {
        var url = matchInfo["url"].StringValue;
        var html = "";
        try
        {
          using (var webClient = new WebClient())
          {
            webClient.Headers["user-agent"] = UserAgent;
            webClient.Encoding = Encoding.UTF8;
            html = webClient.DownloadString(url);
          }
          var regex = matchInfo["regex"].StringValue;
          var match = html.Match(regex);
          var group = decimal.ToInt32(matchInfo["group"].NumberValue.Value);
          if (match.Success)
          {
            if (match.Groups.Count > group - 1)
            {
              var matchedUrl = match.Groups[group].Value;
              downloadUrl(matchedUrl);
            }
            else
            {
              throw new MalwareCheckException($"Regex group not found for {url}, group = {group}");
            }
          }
          else
          {
            throw new MalwareCheckException($"Regex match for {url} failed, regex = {regex}");
          }
        }
        catch (Exception ex)
        when (ex is MalwareCheckException || ex is WebException/* || ex is CryptographicException*/)
        {
          error($"Error fetching {url} {ex.Message}");
          File.WriteAllText(Path.Combine(CheckerOutput, md5(url) + ".html"), html);
        }
      }

    }
    private string md5(string input)
    {
      using (var md5 = MD5.Create())
      {
        var inputBytes = Encoding.UTF8.GetBytes(input);
        var hashBytes = md5.ComputeHash(inputBytes);
        return string.Join("", Enumerable.Range(0, hashBytes.Length).Select(i => hashBytes[i].ToString("x2")));
      }
    }
  }
  class MalwareCheckerModel : NotificationObject
  {

    private string message = "";
    public string Message
    {
      get => message;
      set
      {
        message = value;
        OnPropertyChanged(nameof(Message));
      }
    }

    private double progressValue = 0.0;
    public double ProgressValue
    {
      get => progressValue;
      set
      {
        progressValue = value;
        OnPropertyChanged(nameof(ProgressValue));
      }
    }

    private TaskbarItemProgressState progressState = TaskbarItemProgressState.None;
    public TaskbarItemProgressState ProgressState
    {
      get => progressState;
      set
      {
        progressState = value;
        OnPropertyChanged(nameof(ProgressState));
      }
    }
  }
}
